#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文章捕获和生成系统主程序
功能：
1. 启动mitmproxy进行抓包
2. 定期生成docx文件
3. 监控系统状态
"""

import time
import threading
import signal
import sys
import logging
from pathlib import Path
from datetime import datetime
import subprocess
import yaml

from docx_generator import DocxGenerator

class ArticleCaptureSystem:
    def __init__(self):
        self.config = self.load_config()
        self.setup_logging()
        self.running = False
        self.mitmproxy_process = None
        self.docx_generator = DocxGenerator()
        
    def setup_logging(self):
        """设置日志"""
        log_dir = Path(__file__).parent / 'logs'
        log_dir.mkdir(exist_ok=True)
        
        log_filename = log_dir / f"main_{datetime.now().strftime('%Y%m%d')}.log"
        
        self.logger = logging.getLogger('MainSystem')
        self.logger.setLevel(logging.INFO)
        
        if not self.logger.handlers:
            file_handler = logging.FileHandler(log_filename, encoding='utf-8')
            console_handler = logging.StreamHandler()
            
            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
            file_handler.setFormatter(formatter)
            console_handler.setFormatter(formatter)
            
            self.logger.addHandler(file_handler)
            self.logger.addHandler(console_handler)

    def load_config(self):
        """加载配置文件"""
        config_path = Path(__file__).parent / 'config.yaml'
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            print(f"配置文件未找到: {config_path}")
            return {}
        except Exception as e:
            print(f"加载配置文件失败: {str(e)}")
            return {}

    def start_mitmproxy(self):
        """启动mitmproxy"""
        try:
            # 构建mitmproxy命令
            script_path = Path(__file__).parent / 'addon-cap.py'
            port = self.config.get('mitmproxy', {}).get('port', 8191)
            
            cmd = [
                'mitmdump',
                '-s', str(script_path),
                '-p', str(port),
                '--set', 'confdir=~/.mitmproxy'
            ]
            
            self.logger.info(f"启动mitmproxy: {' '.join(cmd)}")
            
            # 启动mitmproxy进程
            self.mitmproxy_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True
            )
            
            self.logger.info(f"mitmproxy已启动，PID: {self.mitmproxy_process.pid}")
            return True
            
        except Exception as e:
            self.logger.error(f"启动mitmproxy失败: {str(e)}")
            return False

    def stop_mitmproxy(self):
        """停止mitmproxy"""
        if self.mitmproxy_process:
            try:
                self.mitmproxy_process.terminate()
                self.mitmproxy_process.wait(timeout=10)
                self.logger.info("mitmproxy已停止")
            except subprocess.TimeoutExpired:
                self.mitmproxy_process.kill()
                self.logger.warning("强制终止mitmproxy")
            except Exception as e:
                self.logger.error(f"停止mitmproxy时发生错误: {str(e)}")

    def docx_generation_worker(self):
        """docx生成工作线程"""
        while self.running:
            try:
                self.logger.info("开始生成docx文件...")
                self.docx_generator.generate_docx_files()
                
                # 等待下次生成（默认每5分钟检查一次）
                generation_interval = self.config.get('docx_generation', {}).get('interval_minutes', 5)
                time.sleep(generation_interval * 60)
                
            except Exception as e:
                self.logger.error(f"docx生成过程中发生错误: {str(e)}")
                time.sleep(60)  # 发生错误时等待1分钟后重试

    def signal_handler(self, signum, frame):
        """信号处理器"""
        self.logger.info(f"收到信号 {signum}，正在关闭系统...")
        self.stop()

    def start(self):
        """启动系统"""
        self.logger.info("正在启动文章捕获系统...")
        
        # 设置信号处理器
        signal.signal(signal.SIGINT, self.signal_handler)
        signal.signal(signal.SIGTERM, self.signal_handler)
        
        # 启动mitmproxy
        if not self.start_mitmproxy():
            self.logger.error("无法启动mitmproxy，系统退出")
            return
        
        self.running = True
        
        # 启动docx生成线程
        docx_thread = threading.Thread(target=self.docx_generation_worker, daemon=True)
        docx_thread.start()
        
        self.logger.info("文章捕获系统已启动")
        self.logger.info("按 Ctrl+C 停止系统")
        
        try:
            # 主循环
            while self.running:
                time.sleep(1)
                
                # 检查mitmproxy进程状态
                if self.mitmproxy_process and self.mitmproxy_process.poll() is not None:
                    self.logger.error("mitmproxy进程意外退出")
                    break
                    
        except KeyboardInterrupt:
            self.logger.info("收到键盘中断信号")
        finally:
            self.stop()

    def stop(self):
        """停止系统"""
        self.logger.info("正在停止文章捕获系统...")
        self.running = False
        
        # 停止mitmproxy
        self.stop_mitmproxy()
        
        # 最后一次生成docx文件
        try:
            self.logger.info("执行最后一次docx生成...")
            self.docx_generator.generate_docx_files()
        except Exception as e:
            self.logger.error(f"最后一次docx生成失败: {str(e)}")
        
        self.logger.info("文章捕获系统已停止")

def main():
    """主函数"""
    print("文章捕获和生成系统")
    print("=" * 50)
    
    # 检查依赖
    try:
        import mitmproxy
        import pymongo
        import docx
        import yaml
    except ImportError as e:
        print(f"缺少依赖包: {e}")
        print("请安装以下依赖:")
        print("pip install mitmproxy pymongo python-docx pyyaml")
        return
    
    # 检查配置文件
    config_path = Path(__file__).parent / 'config.yaml'
    if not config_path.exists():
        print(f"配置文件不存在: {config_path}")
        print("请创建config.yaml配置文件")
        return
    
    # 启动系统
    system = ArticleCaptureSystem()
    system.start()

if __name__ == "__main__":
    main() 